import argparse
import cv2
import os
from pathlib import Path

def main():
    parser = argparse.ArgumentParser(description='将视频转换为图片序列')
    parser.add_argument('input', help='输入视频文件路径')
    parser.add_argument('-o', '--output', help='输出图片目录路径')
    parser.add_argument('-n', '--interval', type=int, default=1, 
                       help='采样间隔帧数（默认：每帧都保存）')
    
    args = parser.parse_args()
    
    # 检查输入文件是否存在
    if not os.path.isfile(args.input):
        print(f"错误: 输入文件 '{args.input}' 不存在")
        return
    
    # 打开视频文件
    cap = cv2.VideoCapture(args.input)
    if not cap.isOpened():
        print(f"错误: 无法打开视频文件 '{args.input}'")
        return
    
    # 获取视频信息
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    duration = frame_count / fps if fps > 0 else 0
    
    # 打印视频信息
    print("=" * 50)
    print("视频信息:")
    print(f"输入路径: {args.input}")
    print(f"帧率(FPS): {fps:.2f}")
    print(f"总帧数: {frame_count}")
    print(f"分辨率: {width} x {height}")
    print(f"时长: {duration:.2f} 秒")
    print(f"采样间隔: 每 {args.interval} 帧保存一帧")
    print("=" * 50)
    
    if args.output is None:
        return
    
    # 创建输出目录
    os.makedirs(args.output, exist_ok=True)

    # 读取并保存帧
    saved_count = 0
    frame_index = 0
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
            
        if frame_index % args.interval == 0:
            # 生成文件名并保存
            filename = os.path.join(args.output, f"frame_{saved_count:06d}.jpg")
            cv2.imwrite(filename, frame)
            saved_count += 1
            
            # 打印进度
            if saved_count % 50 == 0:
                print(f"已保存 {saved_count} 张图片...")
        
        frame_index += 1
    
    cap.release()
    print(f"转换完成! 共保存 {saved_count} 张图片到 '{args.output}'")

if __name__ == "__main__":
    main()